import pandas as pd
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import recall_score
from sklearn.model_selection import train_test_split

# 加载数据
movies = pd.read_csv('ml-latest-small/movies.csv')
ratings = pd.read_csv('ml-latest-small/ratings.csv')
tags = pd.read_csv('ml-latest-small/tags.csv')

# 预处理电影标签
movie_tags = tags.groupby('movieId')['tag'].apply(lambda x: ' '.join(x)).reset_index()
movies = pd.merge(movies, movie_tags, on='movieId', how='left')
movies['tag'] = movies['tag'].fillna('')

# 分割训练集和测试集
train_ratings, test_ratings = train_test_split(ratings, test_size=0.2, random_state=42)


# 基于内容的推荐（使用电影标签）
def content_based_recommendation():
    tfidf = TfidfVectorizer(stop_words='english')
    tfidf_matrix = tfidf.fit_transform(movies['tag'])
    movie_similarity = cosine_similarity(tfidf_matrix)
    movie_sim_df = pd.DataFrame(movie_similarity, index=movies['movieId'], columns=movies['movieId'])
    return movie_sim_df


content_sim = content_based_recommendation()


def get_content_recommendations(user_id, n=10):
    user_rated = train_ratings[train_ratings['userId'] == user_id]
    if len(user_rated) == 0:
        return []
    top_rated = user_rated.sort_values('rating', ascending=False).head(1)['movieId'].values[0]
    similar_movies = content_sim[top_rated].sort_values(ascending=False).index[1:n + 1]
    return similar_movies.tolist()


# 用户协同过滤
def user_based_collaborative_filtering():
    user_movie_matrix = train_ratings.pivot_table(index='userId', columns='movieId', values='rating').fillna(0)
    user_similarity = cosine_similarity(user_movie_matrix)
    user_sim_df = pd.DataFrame(user_similarity, index=user_movie_matrix.index, columns=user_movie_matrix.index)
    return user_sim_df, user_movie_matrix


user_sim, user_movie_matrix = user_based_collaborative_filtering()


def get_user_based_recommendations(user_id, n=10):
    similar_users = user_sim[user_id].sort_values(ascending=False).index[1:11]
    user_rated = set(train_ratings[train_ratings['userId'] == user_id]['movieId'])
    recommendations = []
    for sim_user in similar_users:
        sim_user_ratings = train_ratings[(train_ratings['userId'] == sim_user) & (train_ratings['rating'] >= 4)]
        for movie in sim_user_ratings['movieId']:
            if movie not in user_rated and movie not in recommendations:
                recommendations.append(movie)
                if len(recommendations) >= n:
                    return recommendations[:n]
    return recommendations[:n]


# 加权混合推荐
def hybrid_recommendation(user_id, n=10, content_weight=0.5, user_weight=0.5, diversity_factor=0.3):
    content_rec = set(get_content_recommendations(user_id, n * 3))
    user_rec = set(get_user_based_recommendations(user_id, n * 3))

    recommendations = {}
    for movie in content_rec:
        recommendations[movie] = recommendations.get(movie, 0) + content_weight
    for movie in user_rec:
        recommendations[movie] = recommendations.get(movie, 0) + user_weight

    movie_genres = movies.set_index('movieId')['genres'].to_dict()
    selected_movies = []
    final_recommendations = []

    sorted_rec = sorted(recommendations.items(), key=lambda x: x[1], reverse=True)

    for movie, score in sorted_rec:
        if len(final_recommendations) >= n:
            break

        if len(selected_movies) > 0:
            current_genres = set(movie_genres.get(movie, "").split('|'))
            max_similarity = 0
            for selected in selected_movies:
                selected_genres = set(movie_genres.get(selected, "").split('|'))
                similarity = len(current_genres & selected_genres) / len(current_genres | selected_genres)
                max_similarity = max(max_similarity, similarity)
            adjusted_score = score * (1 - diversity_factor * max_similarity)
        else:
            adjusted_score = score

        final_recommendations.append((movie, adjusted_score))
        selected_movies.append(movie)

    final_recommendations.sort(key=lambda x: x[1], reverse=True)
    return [x[0] for x in final_recommendations]


# 评估推荐系统
def evaluate_recommendations(recommend_func, test_users=None, n=10):
    if test_users is None:
        test_users = test_ratings['userId'].unique()

    recalls = []
    for user_id in test_users[:100]:
        actual_positives = set(
            test_ratings[(test_ratings['userId'] == user_id) & (test_ratings['rating'] >= 4)]['movieId'])
        if len(actual_positives) == 0:
            continue
        recommended = set(recommend_func(user_id, n))
        if len(recommended) > 0:
            y_true = [1 if movie in actual_positives else 0 for movie in recommended]
            y_pred = [1] * len(y_true)
            recall = recall_score(y_true, y_pred, zero_division=0)
            recalls.append(recall)
    return np.mean(recalls) if recalls else 0


# 获取电影详细信息
def get_movie_details(movie_ids):
    return movies[movies['movieId'].isin(movie_ids)][['movieId', 'title', 'genres']]


# 计算多样性指标
def calculate_diversity(movie_ids):
    genres = []
    for mid in movie_ids:
        genres.extend(movies[movies['movieId'] == mid]['genres'].str.split('|').values[0])
    unique_genres = len(set(genres))
    total_genres = len(genres)
    genre_distribution = {g: genres.count(g) / total_genres for g in set(genres)}
    return unique_genres / total_genres if total_genres > 0 else 0, genre_distribution


# 对比展示各种推荐方法的推荐结果和多样性
def compare_recommendation_methods(user_id, n=10):
    # 获取各种推荐结果
    content_rec = get_content_recommendations(user_id, n)
    user_rec = get_user_based_recommendations(user_id, n)
    hybrid_rec = hybrid_recommendation(user_id, n)

    # 获取电影详情
    content_details = get_movie_details(content_rec)
    user_details = get_movie_details(user_rec)
    hybrid_details = get_movie_details(hybrid_rec)

    # 计算多样性指标
    content_div, content_dist = calculate_diversity(content_rec)
    user_div, user_dist = calculate_diversity(user_rec)
    hybrid_div, hybrid_dist = calculate_diversity(hybrid_rec)

    # 打印结果对比
    print(f"\n=== 推荐方法对比 (用户ID: {user_id}) ===")

    print("\n1. 基于内容的推荐:")
    print(content_details.to_string(index=False))
    print(f"\n多样性指标: {content_div:.2f}")
    print("类型分布:", {k: f"{v:.2f}" for k, v in content_dist.items()})

    print("\n2. 用户协同过滤推荐:")
    print(user_details.to_string(index=False))
    print(f"\n多样性指标: {user_div:.2f}")
    print("类型分布:", {k: f"{v:.2f}" for k, v in user_dist.items()})

    print("\n3. 混合推荐:")
    print(hybrid_details.to_string(index=False))
    print(f"\n多样性指标: {hybrid_div:.2f}")
    print("类型分布:", {k: f"{v:.2f}" for k, v in hybrid_dist.items()})

    # # 绘制多样性对比图
    # try:
    #     import matplotlib.pyplot as plt
    #     methods = ['内容推荐', '用户CF', '混合推荐']
    #     div_scores = [content_div, user_div, hybrid_div]
    #
    #     plt.figure(figsize=(8, 4))
    #     plt.bar(methods, div_scores, color=['skyblue', 'lightgreen', 'salmon'])
    #     plt.title('推荐方法多样性对比')
    #     plt.ylabel('多样性指标')
    #     plt.ylim(0, 1)
    #     for i, v in enumerate(div_scores):
    #         plt.text(i, v + 0.02, f"{v:.2f}", ha='center')
    #     plt.show()
    # except ImportError:
    #     print("\n(注: 要显示多样性对比图，请安装matplotlib库)")


# 评估推荐系统
print("Evaluating recommendation systems...")
print(f"Content-Based Recall: {evaluate_recommendations(get_content_recommendations):.4f}")
print(f"User-Based CF Recall: {evaluate_recommendations(get_user_based_recommendations):.4f}")
print(f"Hybrid Recall: {evaluate_recommendations(hybrid_recommendation):.4f}")

# 对比展示推荐方法
sample_user = test_ratings['userId'].sample(1).values[0]
compare_recommendation_methods(sample_user)